/* Copyright@Han MA
 * 
 * 03/06/2006 
 * 
 * 
 */


using System;
using System.Collections.Generic;
using System.Text;
using System.Drawing;
using System.Drawing.Imaging;

namespace WaterSheds
{

    class Watershed
    {   
        const int INIT = -1;
        const int MASK = -2;
        const int WSHED = 0;
        MyPixel FICTITIOUS = new MyPixel(-1, -1, -1000,-1000,-1000);
        int curlab;
        int curdist;
        Bitmap pic;
        FifoQue que = new FifoQue();
        List<List<MyPixel>> HightSortedList = new List<List<MyPixel>>();

        public Watershed(Bitmap pic)
        {
            for (int i = 0; i < 256; i++)
            {
                HightSortedList.Add(new List<MyPixel>());
            }
            this.pic = new Bitmap(pic);
            unsafe
            {
                BitmapData data = pic.LockBits
                    (new Rectangle(0, 0, pic.Width, pic.Height), ImageLockMode.ReadWrite, PixelFormat.Format24bppRgb);

                byte* imgPtr = (byte*)(data.Scan0);
                byte red, green, blue;
                
                for (int i = 0; i < data.Height; i++)
                {
                    for (int j = 0; j < data.Width; j++)
                    {
                        red = imgPtr[2];
                        green = imgPtr[1];
                        blue = imgPtr[0];                        
                        int hight= (byte)(.299 * red+ .587 * green+ .114 * blue);
                        HightSortedList[hight].Add(new MyPixel(j, i, hight,INIT,0));
                        imgPtr += 3;
                    }
                    imgPtr += data.Stride - data.Width * 3;
                }
            }
            this.curlab = 0;
        }

        public void DoSegment()
        {
            for (int h = 0; h < 256; h++)
            {
                foreach (MyPixel p in HightSortedList[h])
                {
                    p.label = MASK;
                    foreach(MyPixel nighb in Nighbour(p))
                    {
                        if (nighb.label > 0 || nighb.label == WSHED)
                        {
                            p.dist = 1;
                            que.add(p);
                            break;
                        }
                    }
                }
                curdist = -1;
                que.add(FICTITIOUS);
                while (true)
                {
                    MyPixel p = new MyPixel();
                    p = que.remove();
                    if ((p.x == -1) && (p.y == -1))
                    {
                        if (que.IsEmpty())
                            break;
                        else
                        {
                            que.add(FICTITIOUS);
                            curdist++;
                            p = que.remove();
                        }
                    }
                    foreach (MyPixel q in Nighbour(p))
                    {
                        if ((q.dist < curdist) && ((q.label > 0) || (q.label == WSHED)))
                        {
                            if (q.label > 0)
                            {
                                if (p.label == MASK || p.label == WSHED)
                                    p.label = q.label;
                                else
                                {
                                    if (p.label != q.label)
                                        p.label = WSHED;
                                }

                            }
                            else
                            {
                                if (p.label == MASK)
                                    p.label = WSHED;
                            }
                        }
                        else
                        {
                            if ((q.label == MASK) && (q.dist == 0))
                            {
                                q.dist = curdist + 1;
                                que.add(q);
                            }
                        }

                    }
                }
                foreach (MyPixel p in HightSortedList[h])
                {
                    p.dist = 0;
                    if (p.label == MASK)
                    {
                        curlab++;
                        que.add(p);
                        p.label = curlab;
                    }
                    while (!que.IsEmpty())
                    {
                        MyPixel q = new MyPixel();
                        q = que.remove();
                        foreach (MyPixel r in Nighbour(q))
                        {
                            if (r.label == MASK)
                            {
                                que.add(r);
                                r.label = curlab;
                            }
                        }
                    }
                }
            }

            
        }

        public List<MyPixel> Nighbour(MyPixel p) // return 4 connected pixels
        {
            List<MyPixel> temp = new List<MyPixel>();
            if ((p.x-1) >= 0)
            {
                temp.Add(new MyPixel((p.x-1),p.y));
            }

            if ((p.y-1)>=0)
            {
                temp.Add(new MyPixel(p.x,(p.y-1)));
            }
            if ((p.x +1)<=pic.Width )
            {
                temp.Add(new MyPixel((p.x+1),p.y));
            }
            if ((p.y+1)<= pic.Height)
            {
                temp.Add(new MyPixel(p.x,(p.y+1)));
            }

            return temp;
        }

        public Bitmap GetSegPic()
        { 
            
            List<MyPixel> WshdPixels = new List<MyPixel>();
            for (int h=0;h<256;h++)
                foreach (MyPixel p in HightSortedList[h])
                {
                    if (p.label == WSHED)
                        WshdPixels.Add(p);
                }
               unsafe
            {
                BitmapData data = pic.LockBits
                    (new Rectangle(0, 0, pic.Width, pic.Height), ImageLockMode.ReadWrite, PixelFormat.Format24bppRgb);

                byte* imgPtr = (byte*)(data.Scan0);                
                for (int i = 0; i < data.Height; i++)
                {
                    for (int j = 0; j < data.Width; j++)
                    {
                        if (WshdPixels.Contains(new MyPixel(j, i)))
                        {
                            imgPtr[0] = 0;
                            imgPtr[1] = 0;
                            imgPtr[2] = 255;
                        }
                        imgPtr += 3;
                    }
                    imgPtr += data.Stride - data.Width * 3;
                }
            }
            return pic;

        }

    }

}

